/*
 * rim_main.cpp
 *
 *  Created on: Feb 28, 2011
 *      Author: tqlong
 */

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>
#include <armadillo>

#include <ctime>
#include <cassert>

#include <vector>
#include <fstream>
#include <iostream>
#include <map>

#include "entropykmean.h"
#include "common/armavec.h"

namespace po = boost::program_options;
po::options_description desc("Allowed options");
po::variables_map vm;

void process_options(int argc, char** argv);

template <typename vec>
void reginfomax()
{
  Dataset<vec> data;
  data.load(vm["data"].as<std::string>(), arma::arma_ascii);
  data.normalize();
  if (vm.count("norm")) {
//    Normalize<Dataset<vec> > normalize;
//    normalize(data);
//    data.save(vm["data"].as<std::string>()+"-norm");
    PCA<Dataset<vec> > pca;
    pca(data);
  }

  //data.normalize();
  //data.print();

  std::cerr << "n = " << data.n() << " dim = "<< data.dim() << "\n";

  paramset_type p;
  int nc = vm["nc"].as<int>();
  double lambda = vm["lambda"].as<double>();
  int maxiter = vm["iter"].as<int>();
  int memsize = vm["mem"].as<int>();
//  double eta = vm["eta"].as<double>();
  p["lambda"] = lambda/data.n();
  p["iter"] = maxiter;
  p["mem"] = memsize;
//  p.push_back(eta);

  typedef EntropyKMean<vec, Dataset<vec> > Functional;
  Functional F(data, nc);
  F.setParams(p);
  F.optimize();
  F.print();

  typename Functional::vec_dvec_type ps(nc);
  typename Functional::dvec_type pl(nc);
  F.computeProbabilities(F.currentVariable(), ps, pl);

  int n = data.size();
  for (int i = 0; i< n; i++) {
    int maxk = 0;
    for (int k = 1; k< nc; k++)
      if (ps[k][i] > ps[maxk][i]) maxk = k;
    std::cout << "sample " << i << " class " << maxk << "\n";
  }

  std::string mname = vm["mesh"].as<std::string>();
  if (data.dim() == 2) {
    double min = -10, max = 10, step = 0.2;
    std::ofstream fx((mname+".x").c_str()), fy((mname+".y").c_str()),
        fz((mname+".z").c_str()), fc((mname+".class").c_str());
    for (double e1 = min; e1 < max; e1 += step)
      for (double e2 = min; e2 < max; e2 += step) {
        vec x(2);
        x[0] = e1; x[1] = e2;
        double entropy = F.computeEntropy(F.currentVariable(), x);
        int classID = F.computeClass(F.currentVariable(), x);
        fx << x[0] << "\n";
        fy << x[1] << "\n";
        fz << entropy << "\n";
        fc << classID << "\n";
      }
    fx.close(); fy.close(); fz.close(), fc.close();
  }
  std::ofstream fb((mname+".b").c_str());
  for (int k = 0; k < nc; k++)
    fb << F.currentVariable().bias_[k] << " " << arma::norm(F.currentVariable()[k],2) << "\n";
  fb.close();

  std::ofstream fd((mname+".d").c_str());
  for (int i = 0; i < n; i++) {
    for (int d = 0; d < data.dim(); d++)
      fd << data[i].get(d) << " ";
    fd << "\n";
  }
  fd.close();
}

template <typename vec>
void generateData()
{
  int n = vm["n"].as<int>(); // #points in each cluster
  int nc = vm["nc"].as<int>(); // #clusters
  int d = vm["d"].as<int>();
  double r = vm["r"].as<double>();
  Dataset<vec> data;
  for (int c = 0; c < nc; c++) {
    vec center(d);
//    center.randn();
//    center *= 5.0;

    double r = (c+1) * 3;
    center.zeros();
//    center[0] = c * 5.0;
    std::cerr << "c = " << c << " center=\n" << center;
    for (int i = 0; i < n; i++) {
      vec x(d);
      x.randn();
      x /= arma::norm(x,2);
      double sigma = arma::randu(1)[0] * 0.4 + 0.8;
      x *= (r * sigma);
      x += center;
      data.push_back(x);
    }
  }
  data.save(vm["data"].as<std::string>(), arma::arma_ascii);
}

int main(int argc, char** argv) {
  process_options(argc, argv);
  if (vm.count("help") || !vm.count("data")) {
    std::cout << desc << "\n";
    return 0;
  }
  if (vm.count("gen")) {
    if (vm.count("seed"))
      srand(time(NULL));
    generateData <ArmaVec> ();
    return 0;
  }

  boost::posix_time::ptime time_start(boost::posix_time::microsec_clock::local_time());

  reginfomax<ArmaVec> ();

  boost::posix_time::ptime time_end(boost::posix_time::microsec_clock::local_time());
  boost::posix_time::time_duration duration(time_end - time_start);
  std::cerr << "Duration = " << duration << "\n";
  return 0;
}

void process_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce this usage")
    ("data", po::value<std::string>(), "file contains training samples")
    ("norm", "normalize data")
    ("gen", "generate random data")
    ("seed", "using different random seed")
    ("n", po::value<int>()->default_value(10), "number of points")
    ("nc", po::value<int>()->default_value(3), "number of centers")
    ("d", po::value<int>()->default_value(2), "input dimension")
    ("r", po::value<double>()->default_value(1.0), "cluster radius")
    ("lambda", po::value<double>()->default_value(1.0), "regularized parameter")
    ("iter", po::value<int>()->default_value(10), "maximum number of iterations")
    ("mem", po::value<int>()->default_value(4), "LBFGS memory size")
    ("eta", po::value<double>()->default_value(0.5), "learning rate")
    ("mesh", po::value<std::string>()->default_value("mesh"), "mesh files prefix")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}
